package NewCoder;

import java.util.HashMap;
import java.util.Map;

public class BM36 {
    public class Solution {
        Map<TreeNode, Integer> map = new HashMap<>();
        public boolean IsBalanced_Solution(TreeNode root) {
            if (root == null || root.left == null && root.right == null) return true;
            int leftHeight = 0;
            int rightHeight = 0;
            if (map.containsKey(root.left)) {
                leftHeight = map.get(root.left);
            } else {
                leftHeight = depth(root.left);
                map.put(root.left, leftHeight);
            }
            if (map.containsKey(root.right)) {
                rightHeight = map.get(root.right);
            } else {
                rightHeight = depth(root.right);
                map.put(root.right, rightHeight);
            }
            int res = Math.abs(leftHeight - rightHeight);
            if (res > 1) {
                return false;
            }
            return IsBalanced_Solution(root.left) && IsBalanced_Solution(root.right);
        }

        private int depth(TreeNode root) {
            if (root == null) return 0;
            return 1 + Math.max(depth(root.left), depth(root.right));
        }
    }
}
